# Author: WY
# Date: 2023/5/17 16:27

import datetime
import os
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import networkx as nx
import matplotlib.pylab as plt
from pylab import mpl
from sklearn import metrics
import model.GCN_Model as GCN_Model
import model.TCN_Model as TCN_Model

# 防止plot时报错
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
# 设置显示中文字体
mpl.rcParams['font.sans-serif'] = ['SimHei']
# 设置正常显示符号
mpl.rcParams["axes.unicode_minus"] = False
# cuda
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 文件地址
pkl_path = 'result/model.pkl'
val_path = 'result/val.txt'
losses_img_path = 'result/losses.png'
bias_img_path = 'result/bias.png'
variance_img_path = 'result/variance.png'

# 图信息
node_num = 21
edges = [
    (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7),
    (2, 8), (3, 9), (4, 10), (4, 11), (5, 12), (6, 13), (7, 14), (7, 15), (7, 16),
    (11, 17), (12, 18), (13, 19), (15, 20)
]
weights = [
    0.5943, 0.8628, 0.8132, 0.6310, 0.8598, 0.6794, 0.5088,
    0.8638, 0.8536, 0.6448, 0.3974, 0.4254, 0.5618, 0.7458, 0.7001, 0.7193,
    0.6083, 0.7181, 0.7206, 0.7835
]

# 超参数
scopes = 40 # 50
lr = 0.0002 # 0.0002
split = 0.8 # 0.9/0.8
dropout = 0.2
weight_decay = 5e-4
batch_size = 40
pre_step = 10
feature_dim = 13
# 梯度剪裁
clip_grad = False
# GCN
temp_size_conv1 = 13 # 13
temp_size_conv2 = 6 # 6
grad = 2 # 2
gcn_output_size = 3 # 3
# TCN
channels = [50, 50] # [50, 50]
kernel_size = 5 # 5
stride = 1


class GCN_TCN(torch.nn.Module):
    def __init__(self):
        super(GCN_TCN, self).__init__()
        # 图卷积层
        self.GCN_layer = GCN_Model.GCN(
            feature_dim=feature_dim,
            gcn_output_size=gcn_output_size,
            temp_size_conv1=temp_size_conv1,
            temp_size_conv2=temp_size_conv2,
            grad=grad,
            node_num=node_num,
            dropout=dropout
        )
        # 门控
        # self.gate_cell = nn.GRUCell(feature_dim + gcn_output_size, feature_dim + gcn_output_size)
        # 时序卷积层
        self.TCN_layer = TCN_Model.TCN(
            input_size=batch_size,
            output_size=pre_step,
            channels=channels,
            feature_dim=feature_dim + gcn_output_size,
            kernel_size=kernel_size,
            stride=stride,
            dropout=dropout
        )

    def forward(self, input, edge_index, edge_weight):
        # 图卷积提取空间特征
        gcn_out = self.GCN_layer(input, edge_index, edge_weight)
        # print(f"gcn_out: {gcn_out.shape}")

        # 取出主数据部分
        main_data = input[:, 0, :]
        # print(f"main_data: {main_data.shape}")
        # 合并空间张量到主数据的时序特征数据中  [batch, feature_dim + gcn_output_size]
        merge_out = torch.cat([main_data, gcn_out], dim=1)
        # print(f"merge_out: {merge_out.shape}")

        # 时序卷积
        out = self.TCN_layer(merge_out)

        return out


def test(dataloader_train):
    model = GCN_TCN()

    # 通过networkx绘制图,并提取邻接表
    graph = nx.Graph()
    graph.add_nodes_from(np.array(range(node_num)))
    graph.add_edges_from(edges)
    # 制作邻接表和权重表
    edge_index = torch.LongTensor(np.array([edge for edge in graph.edges]).T)
    edge_weight = torch.FloatTensor(weights)

    data = next(iter(dataloader_train))
    # 应GCN的数据类型要求,将 float64(Float) -> float32(Double)
    data = data.to(torch.float32)
    print(f"data: {data.shape}")
    # input与label
    input = data[0: batch_size, :, :]
    print(f"input: {input.shape}")
    label = data[batch_size:, 0, 0].view(1, pre_step)
    print(f"label: {label.shape}")
    # 带入模型
    out = model(input, edge_index, edge_weight)
    print(f"out: {out.shape}")


def train(dataloader_train):
    losses = []
    loss_record = []
    # 实例化模型,损失函数和优化器
    model = GCN_TCN().to(device)
    loss_fun = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)

    # 通过networkx绘制图,并提取邻接表
    graph = nx.Graph()
    graph.add_nodes_from(np.array(range(node_num)))
    graph.add_edges_from(edges)
    # 制作邻接表和权重表
    edge_index = torch.LongTensor(np.array([edge for edge in graph.edges]).T).to(device)
    edge_weight = torch.FloatTensor(weights).to(device)

    # 加载训练结果
    if os.path.exists(pkl_path):
        model.load_state_dict(torch.load(pkl_path))
    # 记录启动时间
    start_date = datetime.datetime.now()
    print(f"start {start_date}")
    # 开始训练
    for scope in range(scopes):
        for batch_idx, (data) in enumerate(dataloader_train):
            # 应GCN的数据类型要求,将 float64(Float) -> float32(Double)
            data = data.to(torch.float32)
            # input与label
            input = data[0:batch_size, :, :].to(device)
            label = data[batch_size:, 0, 0].view(1, pre_step).to(device)
            # 带入模型
            out = model(input, edge_index, edge_weight)
            # loss与backward
            loss = loss_fun(out, label)
            model.zero_grad()
            loss.backward()
            # 梯度剪裁,用于解决梯度爆炸问题
            if clip_grad:
                nn.utils.clip_grad_norm_(model.parameters(), max_norm=1, norm_type=2)

            optimizer.step()

            if np.isnan(loss.item()):
                print("stop")
                return

            if batch_idx % 10 == 0:
                losses.append(loss.item())

        sum = np.sum(losses)
        loss_record.append(sum)
        print(f"scope: {scope}, losses: {sum}")
        losses.clear()

    # 记录结束时间
    end_date = datetime.datetime.now()
    print(f"end {end_date}")
    # 保存模型训练结果
    torch.save(model.state_dict(), pkl_path)
    # 损失函数图像
    fig = plt.figure()
    plt.plot(np.array(range(len(loss_record))), loss_record)
    plt.title(label=f'损失值变化 训练耗时:{end_date - start_date}')
    plt.show()
    # 保存图像
    fig.savefig(losses_img_path)
    # 打印損失值
    np.savetxt(val_path, loss_record, fmt="%.7f", delimiter=',')


def bias_test(dataloader_train):
    # 接受array
    outputs = []
    labels = []
    # 获取模型实例,制作邻接表
    model = GCN_TCN().to(device)

    # 通过networkx绘制图,并提取邻接表
    graph = nx.Graph()
    graph.add_nodes_from(np.array(range(node_num)))
    graph.add_edges_from(edges)
    # 制作邻接表和权重表
    edge_index = torch.LongTensor(np.array([edge for edge in graph.edges]).T).to(device)
    edge_weight = torch.FloatTensor(weights).to(device)

    # 加载训练成果
    if os.path.exists(pkl_path):
        model.load_state_dict(torch.load(pkl_path))
    # 预测结果
    for batch_idx, (data) in enumerate(dataloader_train):
        # 应GCN的数据类型要求,将 float64 -> float32
        data = data.to(torch.float32)
        # input和label
        input = data[0:batch_size, :, :].to(device)
        label = data[batch_size:, 0, 0].view(1, pre_step).tolist()[0]
        # 带入模型,得到结果
        out = model(input, edge_index, edge_weight)
        # 取出预测值
        output = out.cpu().detach().tolist()[0]
        # 记录结果
        outputs.extend(output)
        labels.extend(label)

    # 打印预测结果和label
    xAxis = np.array(range(len(outputs)))
    yAxis_predict = np.array(outputs)
    yAxis_label = np.array(labels)
    fig = plt.figure()
    plt.plot(xAxis, yAxis_label, 'r-.', label='真实值')
    plt.plot(xAxis, yAxis_predict, 'b-.', label='预测值')
    plt.title(label='模型偏差')
    plt.show()
    # 保存图像
    fig.savefig(bias_img_path)


def variance_test(dataloader_test):
    # 接受array
    outputs = []
    labels = []
    # 获取模型实例,制作邻接表
    model = GCN_TCN().to(device)

    # 通过networkx绘制图,并提取邻接表
    graph = nx.Graph()
    graph.add_nodes_from(np.array(range(node_num)))
    graph.add_edges_from(edges)
    # 制作邻接表和权重表
    edge_index = torch.LongTensor(np.array([edge for edge in graph.edges]).T).to(device)
    edge_weight = torch.FloatTensor(weights).to(device)

    # 加载训练成果
    if os.path.exists(pkl_path):
        model.load_state_dict(torch.load(pkl_path))
    # 预测结果
    for batch_idx, (data) in enumerate(dataloader_test):
        # 应GCN的数据类型要求,将 float64 -> float32
        data = data.to(torch.float32)
        # input和label
        input = data[0:batch_size, :, :].to(device)
        label = data[batch_size:, 0, 0].view(1, pre_step).tolist()[0]
        # 带入模型,得到结果
        out = model(input, edge_index, edge_weight)
        # 取出预测值
        output = out.cpu().detach().tolist()[0]
        # 记录结果
        outputs.extend(output)
        labels.extend(label)

    # 打印预测结果和label
    xAxis = np.array(range(len(outputs)))
    yAxis_predict = np.array(outputs)
    yAxis_label = np.array(labels)
    # RMSE
    RMSE = np.sqrt(metrics.mean_squared_error(yAxis_label, yAxis_predict))
    # MAE
    MAE = metrics.mean_absolute_error(yAxis_label, yAxis_predict)
    # 绘图
    fig = plt.figure()
    plt.plot(xAxis, yAxis_label, 'r-.', label='真实值')
    plt.plot(xAxis, yAxis_predict, 'b-.', label='预测值')
    plt.title(label=f"模型方差 RMSE={RMSE} MAE={MAE}")
    plt.legend(loc='upper right')
    plt.show()
    # 保存图像
    fig.savefig(variance_img_path)




